from copy import deepcopy
from functools import partial

from mmengine.config import Config
import numpy as np
import torch

from .build_registry import TRANSFORMS as T


class ImagePreprocessor:
    def __init__(self, cfg, input_shape=(800, 1088)):
        if isinstance(cfg, str):
            cfg = Config.fromfile(cfg)
        self.input_shape = cfg.get('input_shape', input_shape)
        self.transforms = []
        self.pipeline_mapping(cfg)
        self.preprocessor_mapping(cfg)
        self.append('ImageToData')

    def pipeline_mapping(self, cfg):
        cfg = deepcopy(cfg)

        # get test_pipeline config
        test_pipeline = None
        dataset_cfg = cfg.test_dataloader.get('dataset')
        if dataset_cfg:
            test_pipeline = dataset_cfg.get('pipeline')
        if not test_pipeline:
            test_pipeline = cfg.get('test_pipeline')
        assert test_pipeline, 'get test_pipeline config failed!'

        for item in test_pipeline:
            trans_type = item.pop('type')
            if trans_type in set(['LoadAnnotations', 'PackDetInputs']):
                continue

            kwargs = None
            # modify Pad and Resize for fixing the shape of data
            if trans_type == 'Resize':
                kwargs = deepcopy(item)
                if 'scale' in kwargs:
                    kwargs.pop('scale')
                kwargs['input_shape'] = self.input_shape
            elif trans_type == 'Pad':
                kwargs = dict(input_shape=self.input_shape)
                if 'pad_val' in item:
                    kwargs['pad_val'] = item.pad_val.img
            elif trans_type == 'RandomCenterCropPad':
                kwargs = dict(
                    input_shape=self.input_shape,
                    mean=item.mean,
                    std=item.std,
                    test_pad_add_pix=item.test_pad_add_pix
                )

            if not kwargs:
                kwargs = item
            self.append(trans_type, **kwargs)

        self.append('GetImageShape')

    def preprocessor_mapping(self, cfg):
        cfg = deepcopy(cfg)
        preprocessor_cfg = cfg.model.get('data_preprocessor')
        assert preprocessor_cfg, 'get data_preprocessor config failed!'
        assert preprocessor_cfg.type == 'DetDataPreprocessor'

        # ignore unnecessary transform
        preprocessor_cfg.pop('type')
        preprocessor_cfg.pop('batch_augments', None)
        preprocessor_cfg.pop('pad_mask', None)

        # channel exchange
        bgr_to_rgb = preprocessor_cfg.pop('bgr_to_rgb', None)
        rgb_to_bgr = preprocessor_cfg.pop('rgb_to_bgr', None)
        if bgr_to_rgb or rgb_to_bgr:
            self.append('ChannelExchange')

        # normalization
        mean = preprocessor_cfg.pop('mean', None)
        std = preprocessor_cfg.pop('std', None)
        if mean or std:
            if not mean:
                raise Exception(f'Invalid argument "mean": {mean}')
            if not std:
                raise Exception(f'Invalid argument "std": {std}')
            self.append('Normalize', mean=mean, std=std)

        # pad image
        if preprocessor_cfg.pop('pad_size_divisor', None):
            self.append('Pad', input_shape=self.input_shape)

        assert not preprocessor_cfg, f'Unknown transforms: {preprocessor_cfg}'
        self.append('GetBatchInputShape')

    def append(self, name, **func_kwargs):
        # do mapping: transform type --> function
        func = T.get(name, None)
        assert func, f'"{name}" has not been registered yet!'
        self.transforms.append(partial(func, **func_kwargs))

    def preprocess(self, image_path):
        img_meta = {'img_path': image_path}
        img = image_path
        for transform in self.transforms:
            result = transform(img)
            if isinstance(result, (tuple, list)):
                assert len(result) == 2
                img, process_meta = result
                assert isinstance(img, (np.ndarray, torch.tensor))
                assert isinstance(process_meta, dict)
                img_meta.update(process_meta)
            elif isinstance(result, (np.ndarray, torch.tensor)):
                img = result
            elif isinstance(result, dict):
                img_meta.update(result)
            else:
                raise TypeError(
                    f'the return value type of transform `{transform.func}` is '
                    'invalid, expected: Union[np.ndarray | dict | Tuple['
                    f'np.ndarray, dict]], but actually returned: {type(result)}'
                )
        return img, img_meta

    def preprocess_batch(self, batch_images, batch_size=1):
        model_inputs, img_metas = [], []
        for image_path in batch_images:
            img, img_meta = self.preprocess(image_path)
            model_inputs.append(img)
            img_metas.append(img_meta)

        if len(model_inputs) < batch_size:
            pad_num = batch_size - len(model_inputs)
            model_inputs += [model_inputs[-1]] * pad_num
        model_inputs = np.concatenate(model_inputs, axis=0)

        return model_inputs.copy(), img_metas
